#pragma once

#include <memory>
#include <vector>
#include <set>
#include <string>
#include <cxxabi.h>
#include "autograd/autograd.h"
#include "autograd/function.h"
#include "log.h"

namespace  ldl
{
namespace autograd
{
template<typename T>
class Variable
{
public:
    Variable(const T& data, const std::string& name) {
        m_data = std::make_shared<Autograd<T>>(data, name);
    }

    Variable(const T& data) {
        m_data = std::make_shared<Autograd<T>>(data);
    }

    Variable(const T&& data) {
        m_data = std::make_shared<Autograd<T>>(data);
    }

    Variable(std::shared_ptr<Autograd<T>> data) {
        m_data = data;
    }

    virtual ~Variable() = default;

    static std::string demangle_short_name(const char* mangled_name) {
        int status = 0;
        char* demangled = abi::__cxa_demangle(mangled_name, nullptr, nullptr, &status);
        if (status != 0) {
            return mangled_name; // 解修饰失败，返回原始名称
        }

        std::string result(demangled);
        free(demangled);

        // 移除命名空间
        size_t last_colon = result.find_last_of("::");
        if (last_colon != std::string::npos) {
            result = result.substr(last_colon + 1);
        }

        return result;
    }

    void backward() {
        m_data->grad() = std::make_shared<Variable<T>>(ones_like(*(m_data->data())));

        auto cmp = [](std::shared_ptr<Function<T>> a, std::shared_ptr<Function<T>> b) { return a->generation() > b->generation(); };
        std::multiset<std::shared_ptr<Function<T>>, decltype(cmp)> creators(cmp);
        std::set<std::shared_ptr<Function<T>>> seen_creators{};
        seen_creators.insert(m_data->creator());
        creators.insert(m_data->creator());

        while(!creators.empty()) {
            auto node = creators.extract(creators.begin());
            auto creator = node.value();

            creator->backward();

            for(auto item:creator->inputs()) {
                if((item->creator() != nullptr) && (seen_creators.count(item->creator()) <= 0)) {
                    creators.insert(item->creator());
                    seen_creators.insert(item->creator());
                }
            }

            for(auto item:creator->outputs()) {
                item.lock()->grad() = nullptr;
            }
        }
    }

    std::shared_ptr<Autograd<T>> m_data_ptr() const {
        return m_data;
    }

    Variable<T> grad() {
        return *(m_data->grad());
    }

    /**
     * @brief 先将数据暴露出去，以便进行优化
     */
    T& data() {
        return *(m_data->data());
    }

    Variable<T> operator+(Variable<T> other) {
        return Function<T>::add(*this, other);
    }

    Variable<T> operator+=(Variable<T> other) {
        *this = Function<T>::add(*this, other);
        return *this;
    }

    Variable<T> operator+(T other) {
        return Function<T>::add(*this, Variable<T>(other));
    }

    Variable<T> operator-(Variable<T> other) {
        return Function<T>::sub(*this, other);
    }

    Variable<T> operator-=(Variable<T> other) {
        *this = Function<T>::sub(*this, other);
        return *this;
    }

    Variable<T> operator-() {
        return Function<T>::neg(*this);
    }

    Variable<T> operator-(T other) {
        return Function<T>::sub(*this, Variable<T>(other));
    }

    Variable<T> operator*(Variable<T> other) {
        return Function<T>::mul(*this, other);
    }

    Variable<T> operator*(T other) {
        return Function<T>::mul(*this, Variable<T>(other));
    }

    Variable<T> operator/(Variable<T> other) {
        return Function<T>::div(*this, other);
    }

    Variable<T> operator/(T other) {
        return Function<T>::div(*this, Variable<T>(other));
    }

    Variable<T> square() {
        return Function<T>::square(*this);
    }

    // operator T() const {
    //     return *(m_data->data());
    // }

    std::string str() const {
        std::ostringstream oss;
        oss << *m_data;
        return oss.str();
    }

    void cleargrad() {
        m_data->grad() = nullptr;
    }
    friend std::ostream& operator<<(std::ostream& os, const Variable<T>& var)
    {
        return os << *(var.m_data);
    }
    const std::string& name() const {
        return m_data->name();
    }
private:
    std::shared_ptr<Autograd<T>> m_data{};
};

template<typename T>
Variable<T> operator+(T lhs, const Variable<T>& rhs) {
    return Function<T>::add(Variable<T>(lhs), rhs);
}
template<typename T>
Variable<T> operator+(const std::shared_ptr<Autograd<T>>& lhs, const Variable<T>& rhs) {
    return Function<T>::add(Variable<T>(lhs), rhs);
}

template<typename T>
Variable<T> operator-(T lhs, const Variable<T>& rhs) {
    return Function<T>::sub(Variable<T>(lhs), rhs);
}
template<typename T>
Variable<T> operator-(const std::shared_ptr<Autograd<T>>& lhs, const Variable<T>& rhs) {
    return Function<T>::sub(Variable<T>(lhs), rhs);
}

template<typename T>
Variable<T> operator*(T lhs, const Variable<T>& rhs) {
    return Function<T>::mul(Variable<T>(lhs), rhs);
}
template<typename T>
Variable<T> operator*(const std::shared_ptr<Autograd<T>>& lhs, const Variable<T>& rhs) {
    return Function<T>::mul(Variable<T>(lhs), rhs);
}

template<typename T>
Variable<T> operator/(T lhs, const Variable<T>& rhs) {
    return Function<T>::div(Variable<T>(lhs), rhs);
}
template<typename T>
Variable<T> operator/(const std::shared_ptr<Autograd<T>>& lhs, const Variable<T>& rhs) {
    return Function<T>::div(Variable<T>(lhs), rhs);
}
}
}
